# code resource: https://github.com/RRdmlearning/Machine-Learning-From-Scratch/blob/master/logistic_regression/logistic_regression.py

from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np


# sigmoid函数
def sigmoid(x):
    return 1 / (1 + np.exp(x))

class LogisticRegression():
    """
        Params:
            n_iterations: int, 学习的轮数
            learnging_rate: float, 梯度下降的学习率
    """

    def __init__(self, learning_rate = .1, n_iterations = 4000):
        self.learning_rate = learning_rate
        self.n_iterations = n_iterations

    def init_weights(self, n_features):
        # 初始化参数
        # 参数范围(-1/sqrt(N), 1/sqrt(N))
        limit = np.sqrt(1/n_features)
        w = np.random.uniform(-limit, limit, (n_features, 1))
        b = 0
        self.w = np.insert(w, 0, b, axis=0)

    def fit(self, X, y):
        # 训练模型
        m_samples, n_features = X.shape
        self.init_weights(n_features)
        # 为X增加一列特征x1,x1 = 1
        X = np.insert(X, 0, 1, axis=1)
        y = np.reshape(y, (m_samples, 1))

        # 梯度训练n_interations轮
        for i in range(self.n_iterations):
            h_x = X.dot(self.w)
            y_pred = sigmoid(h_x)
            w_grad = X.T.dot(y_pred - y)    # 求梯度
            self.w = self.w - self.learning_rate * w_grad

    def predict(self, X):
        # 预测
        X = np.insert(X, 0, 1, axis=1)
        h_x = X.dot(self.w)
        y_pred = np.round(sigmoid(h_x))
        return y_pred.astype(int)

def main():
    # load dataset
    data = datasets.load_iris()
    X = data.data[data.target != 0]/np.linalg.norm(data.data[data.target != 0], axis=0)
    y = data.target[data.target != 0]
    y[y == 1] = 0
    y[y == 2] = 1

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

    clf = LogisticRegression()
    clf.fit(X_train, y_train)
    y_pred = clf.predict(X_test)
    y_pred = np.reshape(y_pred, y_test.shape)
    accuracy = accuracy_score(y_test, y_pred)
    print("Accuracy:", accuracy)


if __name__ == "__main__":
    main()